#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识库数据导入Neo4j工具
功能：扫描指定路径下的txt文件，将每行数据导入Neo4j图数据库
"""

import os
import re
import jieba
from pathlib import Path
from neo4j import GraphDatabase
from datetime import datetime


class KnowledgeImporter:
    # def __init__(self, uri="bolt://localhost:7687", user="neo4j", password="changsha2025"):
    def __init__(self, uri="bolt://192.168.1.169:7687", user="neo4j", password="Pa@ssw0rd"):

        try:
            self.driver = GraphDatabase.driver(uri, auth=(user, password))
            print(f"✅ 成功连接到Neo4j数据库")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.driver = None
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def clear_database(self):
        """清空数据库"""
        if not self.driver:
            return False
        
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("✅ 数据库已清空")
            return True
    
    def scan_txt_files(self, folder_path):
        """扫描文件夹下的所有txt文件"""
        folder_path = Path(folder_path)
        if not folder_path.exists():
            print(f"❌ 文件夹 {folder_path} 不存在")
            return []
        
        txt_files = list(folder_path.glob("*.txt"))
        print(f"📁 找到 {len(txt_files)} 个txt文件")
        
        all_lines = []
        for txt_file in txt_files:
            try:
                # 尝试UTF-8编码
                with open(txt_file, 'r', encoding='utf-8') as file:
                    lines = file.readlines()
                
                file_lines = []
                for line_num, line in enumerate(lines, 1):
                    line_content = line.strip()
                    if line_content:  # 只处理非空行
                        file_lines.append({
                            'filename': txt_file.name,
                            'filepath': str(txt_file),
                            'line_number': line_num,
                            'content': line_content
                        })
                
                all_lines.extend(file_lines)
                print(f"✅ {txt_file.name}: {len(file_lines)} 行有效数据")
                
            except UnicodeDecodeError:
                # 尝试GBK编码
                try:
                    with open(txt_file, 'r', encoding='gbk') as file:
                        lines = file.readlines()
                    
                    file_lines = []
                    for line_num, line in enumerate(lines, 1):
                        line_content = line.strip()
                        if line_content:
                            file_lines.append({
                                'filename': txt_file.name,
                                'filepath': str(txt_file),
                                'line_number': line_num,
                                'content': line_content
                            })
                    
                    all_lines.extend(file_lines)
                    print(f"✅ {txt_file.name} (GBK): {len(file_lines)} 行有效数据")
                    
                except Exception as e:
                    print(f"❌ 读取文件失败: {txt_file.name} - {e}")
            
            except Exception as e:
                print(f"❌ 读取文件失败: {txt_file.name} - {e}")
        
        print(f"📊 总计读取 {len(all_lines)} 行有效数据")
        return all_lines
    
    def parse_decision_tree_line(self, line_data):
        """解析决策树行数据"""
        content = line_data['content']
        
        # 检查是否是决策树格式 (包含 ->)
        if ' -> ' in content:
            return self._parse_structured_line(line_data)
        else:
            return self._parse_simple_line(line_data)
    
    def _parse_structured_line(self, line_data):
        """解析结构化决策树行"""
        content = line_data['content']
        parts = content.split(' -> ')
        
        # 提取产品编号
        product_match = re.search(r'([A-Z]{1,3}\d{3,4})', parts[0])
        product = product_match.group(1) if product_match else parts[0].strip()
        
        # 提取货号
        cargo_match = re.search(r'\[货号:([^\]]+)\]', content)
        cargo_number = cargo_match.group(1) if cargo_match else product
        
        # 构建节点路径
        nodes = []
        for i, part in enumerate(parts):
            clean_part = re.sub(r'\[货号:[^\]]+\]', '', part).strip()
            if clean_part and i > 0:  # 跳过产品编号部分
                nodes.append({
                    'level': i,
                    'content': clean_part,
                    'type': self._classify_node_type(clean_part, i),
                    'keywords': ' '.join(jieba.cut(clean_part))
                })
        
        return {
            'type': 'decision_tree',
            'product': product,
            'cargo_number': cargo_number,
            'nodes': nodes,
            'filename': line_data['filename'],
            'line_number': line_data['line_number'],
            'original_content': content
        }
    
    def _parse_simple_line(self, line_data):
        """解析简单文本行"""
        content = line_data['content']
        
        # 提取产品编号
        product_matches = re.findall(r'([A-Z]{1,3}\d{3,4})', content)
        product = product_matches[0] if product_matches else 'UNKNOWN'
        
        return {
            'type': 'simple_text',
            'product': product,
            'content': content,
            'keywords': ' '.join(jieba.cut(content)),
            'filename': line_data['filename'],
            'line_number': line_data['line_number'],
            'original_content': content
        }
    
    def _classify_node_type(self, content, level):
        """分类节点类型"""
        if '？' in content or '?' in content:
            return 'question'
        elif '建议' in content or '推荐' in content or '应该' in content:
            return 'solution'
        elif content in ['是', '否', 'Y', 'N', 'yes', 'no']:
            return 'answer'
        elif level == 1:
            return 'problem'
        elif level == 2:
            return 'category'
        elif '则' in content or '如果' in content:
            return 'condition'
        else:
            return 'step'
    
    def import_to_neo4j(self, parsed_data_list):
        """导入数据到Neo4j"""
        if not self.driver:
            print("❌ 数据库连接不可用")
            return False
        
        with self.driver.session() as session:
            imported_count = 0
            
            for data in parsed_data_list:
                try:
                    if data['type'] == 'decision_tree':
                        self._import_decision_tree(session, data)
                    else:
                        self._import_simple_text(session, data)
                    
                    imported_count += 1
                    
                    if imported_count % 100 == 0:
                        print(f"📊 已导入 {imported_count} 条记录...")
                
                except Exception as e:
                    print(f"❌ 导入失败: {data['filename']}:{data['line_number']} - {e}")
            
            print(f"✅ 成功导入 {imported_count} 条记录")
            return True
    
    def _import_decision_tree(self, session, data):
        """导入决策树数据"""
        # 创建或更新产品节点
        session.run("""
            MERGE (p:Product {code: $product})
            SET p.cargo_number = $cargo_number,
                p.name = $product
        """, product=data['product'], cargo_number=data['cargo_number'])
        
        # 创建文件节点
        session.run("""
            MERGE (f:File {name: $filename})
            SET f.path = $filepath
        """, filename=data['filename'], filepath=data.get('filepath', ''))
        
        # 创建行记录节点
        line_id = f"{data['filename']}_{data['line_number']}"
        session.run("""
            MERGE (l:LineRecord {id: $line_id})
            SET l.filename = $filename,
                l.line_number = $line_number,
                l.content = $content,
                l.timestamp = $timestamp
        """, 
        line_id=line_id,
        filename=data['filename'],
        line_number=data['line_number'],
        content=data['original_content'],
        timestamp=datetime.now().isoformat()
        )
        
        # 建立产品到文件的关系
        session.run("""
            MATCH (p:Product {code: $product})
            MATCH (f:File {name: $filename})
            MERGE (p)-[:DOCUMENTED_IN]->(f)
        """, product=data['product'], filename=data['filename'])
        
        # 建立文件到行记录的关系
        session.run("""
            MATCH (f:File {name: $filename})
            MATCH (l:LineRecord {id: $line_id})
            MERGE (f)-[:CONTAINS]->(l)
        """, filename=data['filename'], line_id=line_id)
        
        # 创建决策节点并建立关系
        prev_node_id = f"product_{data['product']}"
        
        for node in data['nodes']:
            node_id = f"{data['product']}_{data['line_number']}_{node['level']}_{hash(node['content'])}"
            
            # 创建决策节点
            session.run("""
                MERGE (n:DecisionNode {id: $node_id})
                SET n.content = $content,
                    n.type = $type,
                    n.level = $level,
                    n.product = $product,
                    n.keywords = $keywords,
                    n.filename = $filename,
                    n.line_number = $line_number
            """,
            node_id=node_id,
            content=node['content'],
            type=node['type'],
            level=node['level'],
            product=data['product'],
            keywords=node['keywords'],
            filename=data['filename'],
            line_number=data['line_number']
            )
            
            # 建立关系
            if node['level'] == 1:
                # 产品到第一个节点
                session.run("""
                    MATCH (p:Product {code: $product})
                    MATCH (n:DecisionNode {id: $node_id})
                    MERGE (p)-[:HAS_ISSUE]->(n)
                """, product=data['product'], node_id=node_id)
            else:
                # 节点之间的关系
                session.run("""
                    MATCH (prev:DecisionNode)
                    WHERE prev.product = $product 
                      AND prev.filename = $filename 
                      AND prev.line_number = $line_number 
                      AND prev.level = $prev_level
                    MATCH (curr:DecisionNode {id: $node_id})
                    MERGE (prev)-[:LEADS_TO]->(curr)
                """, 
                product=data['product'],
                filename=data['filename'],
                line_number=data['line_number'],
                prev_level=node['level'] - 1,
                node_id=node_id
                )
            
            # 建立行记录到决策节点的关系
            session.run("""
                MATCH (l:LineRecord {id: $line_id})
                MATCH (n:DecisionNode {id: $node_id})
                MERGE (l)-[:CONTAINS_NODE]->(n)
            """, line_id=line_id, node_id=node_id)
    
    def _import_simple_text(self, session, data):
        """导入简单文本数据"""
        # 创建或更新产品节点
        if data['product'] != 'UNKNOWN':
            session.run("""
                MERGE (p:Product {code: $product})
                SET p.name = $product
            """, product=data['product'])
        
        # 创建文件节点
        session.run("""
            MERGE (f:File {name: $filename})
        """, filename=data['filename'])
        
        # 创建文本节点
        text_id = f"{data['filename']}_{data['line_number']}"
        session.run("""
            MERGE (t:TextNode {id: $text_id})
            SET t.content = $content,
                t.keywords = $keywords,
                t.filename = $filename,
                t.line_number = $line_number,
                t.product = $product,
                t.timestamp = $timestamp
        """,
        text_id=text_id,
        content=data['content'],
        keywords=data['keywords'],
        filename=data['filename'],
        line_number=data['line_number'],
        product=data['product'],
        timestamp=datetime.now().isoformat()
        )
        
        # 建立关系
        if data['product'] != 'UNKNOWN':
            session.run("""
                MATCH (p:Product {code: $product})
                MATCH (t:TextNode {id: $text_id})
                MERGE (p)-[:MENTIONED_IN]->(t)
            """, product=data['product'], text_id=text_id)
        
        session.run("""
            MATCH (f:File {name: $filename})
            MATCH (t:TextNode {id: $text_id})
            MERGE (f)-[:CONTAINS]->(t)
        """, filename=data['filename'], text_id=text_id)
    
    def create_indexes(self):
        """创建索引以提高查询性能"""
        if not self.driver:
            return
        
        with self.driver.session() as session:
            indexes = [
                "CREATE INDEX product_code_idx IF NOT EXISTS FOR (p:Product) ON (p.code)",
                "CREATE INDEX decision_content_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.content)",
                "CREATE INDEX decision_type_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.type)",
                "CREATE INDEX text_content_idx IF NOT EXISTS FOR (t:TextNode) ON (t.content)",
                "CREATE INDEX file_name_idx IF NOT EXISTS FOR (f:File) ON (f.name)"
            ]
            
            for index_query in indexes:
                try:
                    session.run(index_query)
                    print(f"✅ 索引创建成功")
                except Exception as e:
                    print(f"⚠️  索引创建失败: {e}")


def main():
    """主函数"""
    print("🚀 知识库数据导入Neo4j工具")
    print("=" * 50)
    
    # 配置参数
    folder_path = "/Users/daijunxiong/Downloads/knowledge_txt"
    
    # 创建导入器
    importer = KnowledgeImporter()
    
    if not importer.driver:
        print("❌ 无法连接到Neo4j数据库")
        return
    
    try:
        # 询问是否清空数据库
        choice = input("是否清空现有数据库？(y/N): ").strip().lower()
        if choice == 'y':
            importer.clear_database()
        
        # 扫描txt文件
        print(f"\n📁 扫描文件夹: {folder_path}")
        all_lines = importer.scan_txt_files(folder_path)
        
        if not all_lines:
            print("❌ 没有找到有效数据")
            return
        
        # 解析数据
        print("\n🔄 解析数据...")
        parsed_data = []
        for line_data in all_lines:
            parsed = importer.parse_decision_tree_line(line_data)
            if parsed:
                parsed_data.append(parsed)
        
        print(f"✅ 解析完成，共 {len(parsed_data)} 条有效记录")
        
        # 导入到Neo4j
        print("\n💾 导入数据到Neo4j...")
        success = importer.import_to_neo4j(parsed_data)
        
        if success:
            # 创建索引
            print("\n🔧 创建索引...")
            importer.create_indexes()
            
            print("\n🎉 导入完成！")
            print(f"📊 统计信息:")
            print(f"   - 处理文件数: {len(set(d['filename'] for d in parsed_data))}")
            print(f"   - 导入记录数: {len(parsed_data)}")
            print(f"   - 决策树记录: {len([d for d in parsed_data if d['type'] == 'decision_tree'])}")
            print(f"   - 简单文本记录: {len([d for d in parsed_data if d['type'] == 'simple_text'])}")
            
            print("\n💡 现在您可以使用以下命令进行查询:")
            print("   python neo4j_query_tool.py 'PK511无产物'")
            print("   python neo4j_query_tool.py")
        
    finally:
        importer.close()


if __name__ == "__main__":
    main()
